# -- coding: utf-8 --

import numpy as np
from numpy import linalg as LA
from .config import get_logger

logging = get_logger("sprout.vgg16")


class LoadImgError(RuntimeError):
    pass


class VGGNet:
    def __init__(self):
        self.input_shape = (224, 224, 3)
        self.weight = 'imagenet'
        self.pooling = 'max'
        self._model_vgg = None

    @property
    def model_vgg(self):
        if self._model_vgg is None:
            from tensorflow.keras.applications.vgg16 import VGG16
            self._model_vgg = VGG16(weights=self.weight,
                                    input_shape=(self.input_shape[0], self.input_shape[1], self.input_shape[2]),
                                    pooling=self.pooling,
                                    include_top=False)
            self._model_vgg.predict(np.zeros((1, 224, 224, 3)))
        return self._model_vgg

    def vgg_extract_feat(self, img_path, ignore_error_img=False):
        """
        提取图片特征值。
        img_path : 图片文件位置
        ignore_error_img : 如果发生LoadImgError，是否忽略该异常，仅仅打印日志，并且返回None
        """
        from tensorflow.keras.preprocessing import image
        from tensorflow.keras.applications.vgg16 import preprocess_input as preprocess_input_vgg
        try:
            img = image.load_img(img_path, target_size=(self.input_shape[0], self.input_shape[1]))
        except OSError as e:
            # load_img方法底层调用PIL
            # 如果读取图片异常，则会抛出OSError
            # 捕获该Error，并将其包装为LoadImgError
            # 方便外部处理
            if ignore_error_img:
                logging.error(f"vgg16: load image error: {str(img_path)}")
                return None
            else:
                raise LoadImgError(str(img_path)) from e

        img = image.img_to_array(img)
        img = np.expand_dims(img, axis=0)
        img = preprocess_input_vgg(img)
        feat = self.model_vgg.predict(img)
        norm_feat = feat[0] / LA.norm(feat[0])
        norm_feat = [i.item() for i in norm_feat]
        return norm_feat

